from sqlalchemy import create_engine,String,Integer,Float,Column
from sqlalchemy.orm import declarative_base,sessionmaker

from pydantic import BaseModel,Field

from PIL import Image as PILImage

from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

from langgraph.prebuilt import ToolNode
from langgraph.graph import MessageGraph
from langchain_core.messages import SystemMessage,AIMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import StateGraph,MessagesState,START,END
from langgraph.checkpoint.memory import MemorySaver

import requests
import json
import io

# 创建基类
Base = declarative_base()

# 定义 WeatherInfo 模型
class Weather(Base):
    __tablename__ = "weather"
    city_id = Column(Integer,primary_key=True) #城市id
    city_name = Column(String(50)) #城市名称
    main_weather = Column(String(50)) #主要天气情况
    description = Column(String(100)) #描述
    temperature = Column(Float) #温度
    feels_like = Column(Float) #体感温度
    max_temperature = Column(Float) #最高温度
    min_temperature = Column(Float) #最低温度

# 数据库连接 URI，这里要替换成自己的Mysql 连接信息，以下是各个字段的对应解释：
# root：MySQL 数据库的用户名。
# root：MySQL 数据库的密码。
# localhost：MySQL 服务器的 IP 地址。
# langgraph：要连接的数据库的名称。
# charset=utf8mb4：设置数据库的字符集为 utf8mb4，支持更广泛的 Unicode 字符
# DATABASE_URI = 'mysql+pymysql://root:1qaz2wsx@localhost:3306/langgraph?charset=utf8mb4'
DATABASE_URI = 'mysql+pymysql://root:root@localhost:3306/langgraph?charset=utf8mb4'
engine = create_engine(DATABASE_URI)

# 如果表不存在，则创建表
Base.metadata.create_all(engine)

# 创建会话
Session = sessionmaker(bind=engine)

class WeatherLoc(BaseModel):
    location:str=Field(description="城市的位置名称")

class WeatherInfo(BaseModel):
    """Extracted weather information for a specific city."""
    city_id:int=Field(...,desciption="城市id")
    city_name:str=Field(...,description="城市名称")
    main_weather:str=Field(...,description="主要天气情况")
    description:str=Field(...,description="天气情况的详细描述")
    temperature:float=Field(...,description="当前温度的摄氏度")
    feels_like:float=Field(...,description="当前体感温度的摄氏度")
    max_temperature:float=Field(...,description="最高温度的摄氏度")
    min_temperature:float=Field(...,description="最低温度的摄氏度")

class QueryWeatherSchema(BaseModel):
    """Schema for querying weather information by city name."""
    city_name:str=Field(description="用于查询天气信息的城市名称")

class DeleteWeatherSchema(BaseModel):
    """Schema for deleting weather information by city name."""
    city_name:str=Field(description="用于删除天气信息的城市名称")

## 第一个工具：获取天气信息
@tool(args_schema=WeatherLoc)
def get_weather(location):
    """
        Function to query current weather.
        :param loc: Required parameter, of type string, representing the specific city name for the weather query. \
        Note that for cities in China, the corresponding English city name should be used. For example, to query the weather for Beijing, \
        the loc parameter should be input as 'Beijing'.
        :return: The result of the OpenWeather API query for current weather, with the specific URL request address being: https://api.openweathermap.org/data/2.5/weather. \
        The return type is a JSON-formatted object after parsing, represented as a string, containing all important weather information.
        """
    # Step 1.构建请求
    url = "https://api.openweathermap.org/data/2.5/weather"

    # Step 2.设置查询参数
    params = {
        "q": location,
        "appid": "5c939a7cc59eb8696f4cd77bf75c5a9a",  # 输入API key
        "units": "metric",  # 使用摄氏度而不是华氏度
        "lang": "zh_cn"  # 输出语言为简体中文
    }

    # Step 3.发送GET请求
    response = requests.get(url, params=params)

    # Step 4.解析响应
    data = response.json()
    return json.dumps(data)

## 第二个工具：往数据库插入天气信息
@tool(args_schema=WeatherInfo)
def insert_weather_to_db(city_id,city_name,main_weather,description,temperature,feels_like,max_temperature,min_temperature):
    """Insert weather information into the database."""
    session = Session() #确保每次会话都是新的
    try:
        # 创建天气实例
        weather = Weather(
            city_id=city_id,
            city_name=city_name,
            main_weather=main_weather,
            description=description,
            temperature=temperature,
            feels_like=feels_like,
            max_temperature=max_temperature,
            min_temperature=min_temperature
        )

        # 添加到会话
        session.add(weather)
        # 提交事务
        session.commit()
        return {"messages":[f"天气数据已成功存储至Mysql数据库。"]}

    except Exception as e:
        session.rollback() #出错时数据回滚
        return {"messages": [f"数据存储失败，错误原因：{e}"]}
    finally:
        session.close() #关闭会话

## 第三个工具：从数据库查询天气信息
@tool(args_schema=QueryWeatherSchema)
def query_weather_from_db(city_name):
    """Query weather information from the database by city name.the query city name must be chinese name,for example,
    to query the weather for Beijing,the city_name parameter should be input as '北京'"""
    session = Session()
    try:
        # 查询天气数据
        weather_data = session.query(Weather).filter(Weather.city_name==city_name).first()
        print(weather_data)
        if weather_data:
            return {
                "city_id":weather_data.city_id,
                "city_name":weather_data.city_name,
                "main_weather":weather_data.main_weather,
                "description":weather_data.description,
                "temperature":weather_data.temperature,
                "feels_like":weather_data.feels_like,
                "max_temperature":weather_data.max_temperature,
                "min_temperature":weather_data.min_temperature
            }
        else:
            return {"messages":[f"未找到城市未{city_name}的天气信息"]}
    except Exception as e:
        return {"messages": [f"查询失败，错误原因：{e}"]}
    finally:
        session.close() #关闭会话

## 第四个工具：从数据库删除天气信息
@tool(args_schema=DeleteWeatherSchema)
def delete_weather_from_db(city_name):
    """Delete weather information from the database by city name."""
    session = Session()
    try:
        weather_data = session.query(Weather).filter(Weather.city_name==city_name).first()
        if weather_data:
            session.delete(weather_data)
            session.commit()
            return {"messages":[f"城市为{city_name}的天气已经成功删除"]}
        else:
            return {"messages":[f"未找到城市未{city_name}的天气信息"]}
    except Exception as e:
        session.rollback()
        return {"messages": [f"删除失败，错误原因：{e}"]}
    finally:
        session.close()

tools = [query_weather_from_db,insert_weather_to_db,get_weather,delete_weather_from_db]
tool_node = ToolNode(tools)

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url=api_base,temperature=0)
llm = llm.bind_tools(tools)
print(llm)

def call_model(state):
    print("call_model......")
    messages = state["messages"]
    messages.append(SystemMessage(content="如果数据是从数据库中查询出来的，那么就不进行数据库的插入操作"))
    response = llm.invoke(messages)
    return {"messages":[response]}

def should_continue(state):
    print("should_continue...")
    messages = state["messages"]
    last_message = messages[-1]
    print(last_message)
    print(last_message.tool_calls)
    if not last_message.tool_calls:
        return "end"
    elif last_message.tool_calls[0]["name"] == "delete_weather_from_db":
        return "run_tool"
    else:
        return "continue"

def run_tool(state):
    print("run_tool.....")
    new_messages = []
    messages = state["messages"]
    tool_calls = messages[-1].tool_calls

    tools = [delete_weather_from_db]
    tools = {t.name:t for t in tools}

    for tool_call in tool_calls:
        tool = tools[tool_call["name"]]
        result = tool.invoke(tool_call["args"])
        new_messages.append({
            "role":"tool",
            "name":tool_call["name"],
            "content":result,
            "tool_call_id":tool_call["id"]
        })
    print(new_messages)
    return {"messages":new_messages}

workflow = StateGraph(MessagesState)

workflow.add_node("agent",call_model)
workflow.add_node("run_tool",run_tool)
workflow.add_node("action",tool_node)

workflow.add_edge(START,"agent")
workflow.add_conditional_edges("agent",should_continue,
                               {
                                   "continue":"action",
                                   "end":END,
                                   "run_tool":"run_tool"
                               })

workflow.add_edge("action","agent")
workflow.add_edge("run_tool","agent")

memory = MemorySaver()
graph = workflow.compile(checkpointer=memory,interrupt_before=["run_tool"])

def display(image_data):
    image = PILImage.open(io.BytesIO(image_data))
    image.show()

# display(graph.get_graph().draw_mermaid_png())

config = {"configurable":{"thread_id":"201"}}

# for chunk in graph.stream({"messages":"天津的天气怎么样?并把数据插入数据库中"},config=config,stream_mode="values"):
#     chunk["messages"][-1].pretty_print()

# for chunk in graph.stream({"messages":"帮我同时查询一下北京和杭州的天气情况，比较哪个城市更适合出游，并且把数据插入数据库中"},config=config,stream_mode="values"):
#     chunk["messages"][-1].pretty_print()

# for chunk in graph.stream({"messages":"帮我删除数据库中的北京的天气信息"},config=config,stream_mode="values"):
#     chunk["messages"][-1].pretty_print()

# for chunk in graph.stream({"messages":"帮我查询数据库中的天津的天气信息"},config=config,stream_mode="values"):
#     chunk["messages"][-1].pretty_print()

# for chunk in graph.stream({"messages":"帮我删除数据库中上海的天气信息"},config=config,stream_mode="values"):
#     state = graph.get_state(config=config)
#
#     # 检查是否有任务，如果没有则结束循环
#     if not state.tasks:
#         chunk["messages"][-1].pretty_print()
#         break
#
#     if state.tasks[0].name == "run_tool":
#         while True:
#             user_input = input("是否允许执行删除操作？请输入'是'或'否'：")
#             if user_input in ["是","否"]:
#                 break
#             else:
#                 print("输入错误，请输入'是'或'否'。")
#
#         if user_input == "是":
#             graph.update_state(config=config, values=chunk)
#             for event in graph.stream(None, config, stream_mode="values"):
#                 event["messages"][-1].pretty_print()
#         elif user_input == "否":
#             state = graph.get_state(config=config)
#             tool_call_id = state.values["messages"][-1].tool_calls[0]["id"]
#             print(tool_call_id)
#
#             # 我们现在需要构造一个替换工具调用。把参数改为“xxsd”，请注意，我们可以更改任意数量的参数或工具名称-它必须是一个有效的
#             new_messages = {
#                 "role":"tool",
#                 "name":"delete_weather_from_db",
#                 # 这是得到的用户不允许操作的反馈
#                 "content":"管理员不允许执行该操作！",
#                 "tool_call_id":tool_call_id
#             }
#             graph.update_state(config,{"messages":[new_messages]},as_node="run_tool")
#             for event in graph.stream(None,config=config,stream_mode="values"):
#                 event["messages"][-1].pretty_print()

def run_multi_round_dialogue(graph,config):
    while True:
        # 询问用户输入操作，允许退出
        user_input = input("请输入您的问题（例如：'帮我查询上海的天气数据'），输入'退出'结束对话：")

        if user_input.lower() == "退出":
            print("对话已结束。")
            break

        # 启动对话，根据用户的输入进行处理
        for chunk in graph.stream({"messages":user_input},config,stream_mode="values"):
            state = graph.get_state(config)

            if not state.tasks:
                if "messages" in chunk and len(chunk["messages"])>0:
                    print("人工智能助理",chunk["messages"][-1].content)
                break

            # 处理动态断点的任务
            if state.tasks[0].name == "run_tool":
                user_approval = None
                while True:
                    user_approval = input("是否允许执行删除操作？请输入'是'或'否'：")
                    if user_approval in ["是", "否"]:
                        break
                    else:
                        print("输入错误，请输入'是'或'否'。")

                if user_approval == "是":
                    graph.update_state(config=config,values=chunk)
                    for event in graph.stream(None,config,stream_mode="values"):
                        if "messages" in event and len(event["messages"]) > 0:
                            print("人工智能助理：", event["messages"][-1].content)
                elif user_approval == "否":
                    state = graph.get_state(config)
                    tool_call_id = state.values["messages"][-1].tool_calls[0]["id"]

                    # 构造一个反馈消息来停止操作
                    new_message = {
                        "role":"tool",
                        "content":"管理员不允许执行该操作！",
                        "name":"delete_weather_from_db",
                        "tool_call_id":tool_call_id
                    }
                    graph.update_state(config,{"messages":[new_message]},as_node="run_tool")
                    for event in graph.stream(None, config, stream_mode="values"):
                        if "messages" in event and len(event["messages"]) > 0:
                            print("人工智能助理：", event["messages"][-1].content)


run_multi_round_dialogue(graph,config)

# weather = get_weather(location="chongqing")
# print(weather)

# resp = call_model({"messages":"重庆的天气怎么样"})
# print(resp)